import easyocr
from flask import Flask, request
from vosk import Model, KaldiRecognizer, SetLogLevel
from flask_cors import *
import json
import wave

app = Flask(__name__)
CORS(app, resources=r'/*')

# 支持的文件格式
app.config['ALLOWED_EXTENSIONS'] = {'png', 'jpg'}  # 集合类型


# 判断文件名是否是我们支持的格式
def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1] in app.config['ALLOWED_EXTENSIONS']


@app.route('/uploadOcr', methods=['POST'])
def uploadOcr():
    upload_file = request.files['file']
    if allowed_file(upload_file.filename):
        reader = easyocr.Reader(['ch_sim', 'en'], model_storage_directory='easyocr-model')
        result = reader.readtext(upload_file.stream.read())
        sen = ""
        for i in result:
            word = i[1]
            sen = sen + word
        print(sen)
        if sen == '':
            sen = "识别失败，请重新上传清晰的图片！"
        return sen
    else:
        print('请上传正确的图片格式！')
        return '识别失败，请上传正确的图片格式！'


@app.route('/uploadAsr', methods=['POST'])
def uploadAsr():
    upload_file = request.files['file']
    model = Model("vosk-model-small-cn-0.22")
    SetLogLevel(0)
    wf = wave.open(upload_file.stream, "rb")
    if wf.getnchannels() != 1 or wf.getsampwidth() != 2 or wf.getcomptype() != "NONE":
        print("Audio file must be WAV format mono PCM.")
        exit(1)
    rec = KaldiRecognizer(model, wf.getframerate())
    rec.SetWords(True)
    res = json.loads(rec.Result())
    sen = str(res['text'])
    while True:
        data = wf.readframes(4000)
        if len(data) == 0:
            break
        if rec.AcceptWaveform(data):
            res = json.loads(rec.Result())
            sen = sen + str(res['text']).replace(' ', '') + ' '
        else:
            rec.PartialResult()
    res = json.loads(rec.FinalResult())
    sen = sen + str(res['text']).replace(' ', '')
    if sen == '':
        sen = "什么也没有听到！"
    print('识别结果：' + sen)
    return sen


if __name__ == '__main__':
    app.run(port=5000, debug=True)
